from PIL import Image
import cv2
import numpy as np
import onnxruntime
from utils.utils import image_transform, pil_to_cv

"""
该py文件为测试文件夹，可以将对resources文件夹下的img文件夹中的图片匹配到face_datas文件中的图片
"""

model_path = "resources/model.onnx"
face_feature_path = "resources/face_feature_vector.txt"
image_paths = './resources/images_path.txt'


def onnx_runtime(img):
    ort_session = onnxruntime.InferenceSession(model_path)
    ort_input1 = {'input': img}
    ort_output1 = ort_session.run(['output'], ort_input1)[0]

    return ort_output1


def face_image_matching(img):
    face_features = np.loadtxt(face_feature_path)
    images_path = []

    with open(image_paths, 'r') as f:
        for line in f.readlines():
            images_path.append(line.split('\n')[0])
    f.close()

    features_distance = []
    for i in range(face_features.shape[0]):
        vector = face_features[i, :]
        vector.reshape(1, 128)
        features_distance.append(np.linalg.norm(img - vector, axis=1)[0])

    print(features_distance)
    if np.min(np.array(features_distance)) < 0.85:
        min_arg = np.argmin(np.array(features_distance))
        mathing = cv2.imread(images_path[min_arg])
        cv2.imshow('math_image', mathing)
        cv2.waitKey(0)
    else:
        print('数据库中无匹配的人脸信息')


while True:
    image_path = input('Input image filename:')

    try:
        image = Image.open(image_path)
        image = pil_to_cv(image)
        cv2.imshow('initial_image', image)
    except:
        print('Image Open Error! Try again!')
        continue

    image = image_transform(image)
    output1 = onnx_runtime(image)
    face_image_matching(output1)
